#include "luster_vm.hh"

#include <assert.h>

#include "luster_log.hh"
#include "luster_test.hh"
#include "luster_vec.hh"

namespace luster {

typedef U32 (*handle_op_func)(Vm *vm, S32 value);

static U32 op_HALT(Vm *vm, S32 value) {
	vm->is_stop = true;
	return vm->pc + 1;
}

static U32 op_CALL(Vm *vm, S32 value) {
	vec_append(&vm->cstack,
		CallFrame{
			.return_pc = vm->pc + 1,
			.on_call_vstack_len = static_cast<U32>(vm->vstack.len),
		});
	return value;
}

static U32 op_RETURN(Vm *vm, S32 value) {
	vm->cstack.len -= 1;

	auto const &cf = vm->cstack[vm->cstack.len];

	// RETURN 时的 vstack 的栈顶是要返回的值。将返回值压入 CALL 时 vstack
	// 的栈顶。
	//
	// 只要满足调用约定，RETURN 时的 vm->vstack.len 一定大于 CALL 时栈顶，
	// 所以不必调用 vec_append 来防止栈越界，只需要直接设置栈的长度。
	vm->vstack[cf.on_call_vstack_len] = vm->vstack[vm->vstack.len - 1];
	vm->vstack.len = cf.on_call_vstack_len + 1;

	return cf.return_pc;
}

static U32 op_PUSH_U32(Vm *vm, S32 value) {
	vec_append_zero(&vm->vstack)->as.u32 = value;
	return vm->pc + 1;
}

static U32 op_POP(Vm *vm, S32 value) {
	vm->vstack.len -= 1;
	return vm->pc + 1;
}

static VmValue &prev(Vm *vm) {
	return vm->vstack.data[vm->vstack.len - 2];
}

static VmValue &top(Vm *vm) {
	return vm->vstack.data[vm->vstack.len - 1];
}

static VmValue &next(Vm *vm) {
	return vm->vstack.data[vm->vstack.len];
}

static U32 op_TUCK(Vm *vm, S32 value) {
	auto a = prev(vm).as.u32;
	auto b = top(vm).as.u32;

	prev(vm).as.u32 = b;
	top(vm).as.u32 = a;
	next(vm).as.u32 = b;

	vm->vstack.len += 1;
	return vm->pc + 1;
}

static U32 op_ADD(Vm *vm, S32 value) {
	vm->vstack.len -= 1;
	top(vm).as.u32 += next(vm).as.u32;
	return vm->pc + 1;
}

static U32 op_SUB(Vm *vm, S32 value) {
	vm->vstack.len -= 1;
	top(vm).as.u32 -= next(vm).as.u32;
	return vm->pc + 1;
}

static U32 op_MUL(Vm *vm, S32 value) {
	vm->vstack.len -= 1;
	top(vm).as.u32 *= next(vm).as.u32;
	return vm->pc + 1;
}

static U32 op_DIV(Vm *vm, S32 value) {
	vm->vstack.len -= 1;
	top(vm).as.u32 /= next(vm).as.u32;
	return vm->pc + 1;
}

static U32 op_MOD(Vm *vm, S32 value) {
	vm->vstack.len -= 1;
	top(vm).as.u32 %= next(vm).as.u32;
	return vm->pc + 1;
}

static U32 op_JUMP(Vm *vm, S32 value) {
	return value;
}

static U32 op_JUMP_ZERO(Vm *vm, S32 value) {
	if (vm->vstack[vm->vstack.len - 1].as.u32 == 0) {
		return value;
	} else {
		return vm->pc + 1;
	}
}

static U32 op_JUMP_NON_ZERO(Vm *vm, S32 value) {
	if (vm->vstack[vm->vstack.len - 1].as.u32 != 0) {
		return value;
	} else {
		return vm->pc + 1;
	}
}

static handle_op_func HANDLE[]{
#define _(CODE) op_##CODE,
	luster_VMCODE(_)
#undef _
};

static char const *INST_NAME[]{
#define _(CODE) #CODE,
	luster_VMCODE(_)
#undef _
};

void vm_run(Vm *vm, Inst const *inst) {
	vm->pc = 0;
	vm->inst = inst;
	vm->is_stop = false;

	vm->cstack.len = 0;
	vm->vstack.len = 0;

	while (!vm->is_stop) {
		auto inst = vm->inst[vm->pc];

		LOG(DEBUG) << "vm: " << INST_NAME[(UInt)inst.code] << LOGEND;
		vm->pc = HANDLE[(U8)inst.code](vm, inst.value);
	}
}

luster_TEST(vm_basic) {
	Vm vm;

	Inst inst[] = {
		{OpCode::PUSH_U32, 20},
		{OpCode::PUSH_U32, 22},
		{OpCode::ADD, 0},
		{OpCode::HALT, 0},
	};

	vm_run(&vm, inst);

	assert(vm.vstack.data[0].as.u32 == 42);
}

luster_TEST(vm_call) {
	Vm vm;
	Vec<Inst> insts;

	using C = OpCode;

	// 第一条指令是跳转调用主函数的指令，但是目前还没有定义好主函数，
	// 所以跳转地址先填零。
	//
	// 主函数返回之后调用 HALT 指令结束程序。
	auto call_main = vec_append(&insts, {C::CALL, 0});
	vec_append(&insts, {C::HALT});

	// 函数 add3 接受 3 个参数，返回它们的和。

	S32 fn_add3 = insts.len;

	vec_append(&insts, {C::ADD, 0});
	vec_append(&insts, {C::ADD, 0});
	vec_append(&insts, {C::RETURN});

	// 主函数。

	call_main->value = insts.len;

	// 计算 20 + 20 + 2
	vec_append(&insts, {C::PUSH_U32, 20});
	vec_append(&insts, {C::PUSH_U32, 20});
	vec_append(&insts, {C::PUSH_U32, 2});
	vec_append(&insts, {C::CALL, fn_add3});

	// 计算上一个求和的结果再加 2 加 3
	vec_append(&insts, {C::PUSH_U32, 2});
	vec_append(&insts, {C::PUSH_U32, 3});
	vec_append(&insts, {C::CALL, fn_add3});

	// 主函数返回
	vec_append(&insts, {C::RETURN});

	vm_run(&vm, insts.data);
	assert(vm.vstack.data[0].as.u32 == 47);
}

luster_TEST(vm_gcd) {
	Vm vm;
	Vec<Inst> insts;

	using C = OpCode;

	auto call_main = vec_append(&insts, {C::CALL});
	vec_append(&insts, {C::HALT});

	//
	// gcd
	//
	auto fn_gcd = insts.len;

	// [a, b]
	auto return_gcd = vec_append(&insts, {C::JUMP_NON_ZERO});

	// [a, b]
	vec_append(&insts, {C::POP});
	// [a]
	vec_append(&insts, {C::RETURN});

	return_gcd->value = insts.len;

	// [a, b]
	vec_append(&insts, {C::TUCK});
	// [b, a, b]
	vec_append(&insts, {C::MOD});
	// [b, a%b]
	vec_append(&insts, {C::CALL, (S32)fn_gcd});
	// [gcd(a, b)]
	vec_append(&insts, {C::RETURN});

	//
	// main
	//
	call_main->value = (S32)insts.len;

	vec_append(&insts, {C::PUSH_U32, 2097151}); // 7 7 127 337
	vec_append(&insts, {C::PUSH_U32, 674}); // 337 2
	vec_append(&insts, {C::CALL, (S32)fn_gcd});
	vec_append(&insts, {C::RETURN});

	vm_run(&vm, insts.data);

	assert(top(&vm).as.u32 == 337);
}

} // namespace luster
